#!/usr/bin/perl -wc

package RFCMail::DNS::AsyncDNSUtils;

use strict;
use base qw(Exporter);
use Carp;
use Errno;
use IO::Socket::INET;
use Event;
use Net::DNS;
use RFCMail::Util::Callback qw(callback);


BEGIN {
    our @EXPORT_OK = qw(async_dns_query
			async_dns_multi_query
			async_dns_query_udp
			async_dns_query_tcp
			async_dns_query_existing_tcp);
    our $RES = Net::DNS::Resolver->new;
    our $UDP_BUFFER_SIZE = 1500; # technically, the payload shouldn't exceed 512
    our (%Tracker, $Udp_Socket, $Udp_Watcher); # tracking for UDP queries
}


sub async_dns_query {
    my ($callback, @query) = @_;
    async_dns_query_udp(sub {
	my $result = shift;
	if ($result->header->tc) {
	    # Response was truncated; use TCP instead
	    async_dns_query_tcp($callback, @query);
	}
	else {
	    callback($callback, $result);
	}
    }, @query);
}


sub async_dns_multi_query {
    my ($callback, @queries) = @_;
    my $outstanding = scalar(@queries);
    my @results;
    for (@queries) {
	async_dns_query(sub {
	    push @results, shift;
	    --$outstanding
		or callback($callback, @results);
	}, @$_);
    }
}


sub async_dns_query_udp {
    my $callback = shift;
    our ($RES);
    my $query = Net::DNS::Packet->new(@_)
	or croak("unable to generate query packet");
    my $sender = Event->timer(after => 0, interval => $RES->retrans, parked => 1);
    gazette($query, $callback, $sender);
    $sender->cb([new_udp_query($query), 'transmit_query_handler']);
    $sender->start;
}


sub async_dns_query_tcp {
    my $callback = shift;
    our ($RES);
    my ($ns) = $RES->nameservers;
    my $sock = IO::Socket::INET->new(PeerAddr => $ns, PeerPort => $RES->port,
				     Proto => "tcp", Blocking => 0);
    async_dns_query_existing_tcp(sub {
	$sock->close;
	callback($callback, @_);
    }, $sock, @_);
}


sub async_dns_query_existing_tcp {
    my $callback = shift;
    my $sock = shift;
    our ($RES);
    my $packet = Net::DNS::Packet->new(@_)
	or croak("unable to generate query packet");
    my $data = pack("n/a*", $packet->data);
    # done with $packet for now, but we'll assemble the response packet in it later
    undef $packet;
    my $event;
    Event->io(fd => $sock, poll => 'w', timeout => $RES->tcp_timeout, cb => sub {
	$event = shift;
	if ($event->got) {
	    # ready to write
	    my $wrote = $sock->syswrite($data);
	    if (defined($wrote)) {
		# cut the bit we sent from $data
		substr($data, 0, $wrote, "");
		# return if more to send
		return
		    if $data;
	    }
	    else {
		# some kind of error; just return if it's a blocking issue
		return
		    if $!{EAGAIN} || $!{EWOULDBLOCK};
	    }
	}
	# reaching this point means a timeout, a fatal error, or successful transmission
	# no more write-watcher events in any case
	$event->w->cancel;
	# untransmitted data means we failed
	if ($data) {
	    callback($callback, undef);
	    return;
	}
	# success, on the other hand, means we read back into the now-empty $data
	Event->io(fd => $sock, poll => 'r', timeout => $RES->tcp_timeout, cb => sub {
	    $event = shift;
	    if ($event->got) {
		# ready to read; compute the remaining bytes wanted
		# note that we first want two bytes to determine the size, then that many more
		my $have = length($data);
		my $want = ($have < 2 ? 2 : unpack("n", $data) + 2) - $have;
		my $got = $sock->sysread($data, $want, $have);
		if (defined($got)) {
		    # some data has been appended to $data; return if not complete
		    return
			if length($data) < 2;
		    return
			if length($data) < unpack("n", $data) + 2;
		    # we seem to have a winner; extract the packet
		    $data = unpack("n/a", $data);
		    $packet = Net::DNS::Packet->new(\$data);
		    undef $data;
		}
		else {
		    # some kind of error; just return if it's a blocking issue
		    return
			if $!{EAGAIN} || $!{EWOULDBLOCK};
		}
	    }
	    # reaching this point means a timeout, a fatal error, or successful reception
	    # these are all terminal cases, and $packet is correct for all cases
	    $event->w->cancel;
	    callback($callback, $packet);
	});
    });
}


# From here on down it's internal subs only. Not intended for external use.


# gazette() and finalise() manage the list of outstanding UDP queries, matching requests to responses.

# To gazette is to publish a fact in an official manner
# Outstanding callbacks are gazetted in %Tracker, indexed by query ID
# NB: The query ID will be adjusted if needs be to avoid a clash
sub gazette {
    my ($query, $callback, $sender) = @_;
    our (%Tracker, $Udp_Socket, $Udp_Watcher);
    my $id = $query->header->id;
    while (exists($Tracker{$id})) { $id++; $id %= 0x10000; }
    $query->header->id($id);
    if (not scalar(%Tracker)) {
	# This will be the first entry in the tracker; create the socket and watcher
	$Udp_Socket = IO::Socket::INET->new(Proto => 'udp', Blocking => 0)
	    or die "unable to open socket for DNS queries ($!)";
	$Udp_Watcher = Event->io(fd => $Udp_Socket, cb => \&read_response_handler);
    }
    $Tracker{$id} = [$callback, $sender];
}

# finalise() issues the callback for the given query ID (if it exists) and delists the callback
# The UDP socket and its watcher are also decommissioned if there is nothing else to track
sub finalise {
    my ($id, $response) = @_;
    our (%Tracker, $Udp_Socket, $Udp_Watcher);
    die "attempt to finalise unknown query (ID $id)"
	if not exists($Tracker{$id});
    my ($callback, $sender) = @{delete($Tracker{$id})};
    $sender->cancel;
    if (not scalar(%Tracker)) {
	# The last tracker entry was just deleted; close the socket and cancel the watcher
	$Udp_Watcher->cancel;
	undef $Udp_Watcher;
	$Udp_Socket->close;
	undef $Udp_Socket;
    }
    callback($callback, $response);
}


# These methods handle transmission of a specific query-in-progress

sub new_udp_query {
    my ($query) = @_;
    our ($RES);
    return (bless [$query->header->id, $query->data, $RES->retry, 0], __PACKAGE__);
}

sub id { $_[0]->[0] }

sub data { $_[0]->[1] }

sub keep_trying { $_[0]->[2] > $_[0]->[3] }

sub count_retrans { $_[0]->[3]++ }

sub transmit_query_handler {
    my ($self) = @_;
    our ($RES, $Udp_Socket);
    if ($self->keep_trying) {
	my $count = $self->count_retrans;
	# cycle through nameservers on successive queries
	my @ns = $RES->nameservers;
	my $ns = $ns[$count % scalar(@ns)];
	my $ns_addr = pack_sockaddr_in($RES->port, inet_aton($ns));
	# attempt a send; if it works, it works
	$Udp_Socket->send($self->data, MSG_DONTWAIT, $ns_addr);
    }
    else {
	# Query is outstanding, but we're out of bullets, so call it a timeout.
	finalise($self->id, undef);
    }
}


# This is the handler for all incoming UDP responses.

sub read_response_handler {
    my $event = shift;
    my ($data, $answer);
    our ($UDP_BUFFER_SIZE);
    $event->w->fd->recv($data, $UDP_BUFFER_SIZE, MSG_DONTWAIT);
    return
	if $!{EAGAIN} || $!{EWOULDBLOCK};
    # abort if data doesn't make sense as a DNS message
    $answer = Net::DNS::Packet->new(\$data)
	or return;
    # otherwise, we have a useful answer
    finalise($answer->header->id, $answer);
}

1;
